package code.classic150on2025;

import java.util.*;

public class TheHeap {
    // 502. IPO
    public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
        // affordable表示负担得起的项目，比较规则为：利润 由高到低
        PriorityQueue<int[]> affordable = new PriorityQueue<>((o1, o2) -> Integer.compare(o2[0], o1[0]));
        // not表示负担不起的项目，比较规则为：成本 由低到高
        PriorityQueue<int[]> not = new PriorityQueue<>((o1, o2) -> Integer.compare(o1[1], o2[1]));
        // 初始化
        int n = profits.length;
        for (int i = 0; i < n; i++) {
            int[] add = new int[]{profits[i], capital[i]};
            if (capital[i] <= w)
                affordable.add(add);
            else
                not.add(add);
        }
        // 当affordable不为空时，且总项目少于k时，循环，找到affordable里赚的最多的选取，然后更新not和affordable
        int count = 0;
        while (!affordable.isEmpty() && count < k) {
            int[] now = affordable.poll();
            w += now[0];
            count++;
            while (!not.isEmpty() && w >= not.peek()[1]) {
                int[] trans = not.poll();
                affordable.add(trans);
            }
        }
        return w;
    }

    // 373. 查找和最小的 K 对数字
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> lists = new ArrayList<>();
        int n1 = nums1.length, n2 = nums2.length;
        PriorityQueue<int[]> pq = new PriorityQueue<>((o1, o2) -> Integer.compare(o1[0], o2[0]));
        for (int i = 0; i < n1; i++) {
            pq.add(new int[]{nums1[i] + nums2[0], i, 0}); // 「和、nums1下标、nums2下标」
        }
        while (lists.size() < k) {
            int[] now = pq.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[now[1]]);
            list.add(nums2[now[2]]);
            lists.add(list);
            if (now[2] + 1 < n2)
                pq.add(new int[]{nums1[now[1]] + nums2[now[2] + 1], now[1], now[2] + 1}); // 只将(now[0], now[1] + 1)入队，不将(now[0] + 1, now[1])入队
        }
        return lists;
    }

    // 215. 数组中的第K个最大元素
    public int findKthLargest(int[] nums, int k) {
        List<Integer> list = new ArrayList<>();
        for (int num : nums)
            list.add(num);
        return findKth(list, k);
    }

    private int findKth(List<Integer> list, int k) {
        // 随机选择一个基数
        Random random = new Random();
        int pivot = list.get(random.nextInt(list.size()));
        // 根据三个数组递归🌟三路快排
        List<Integer> smaller = new ArrayList<>();
        List<Integer> equal = new ArrayList<>();
        List<Integer> bigger = new ArrayList<>();
        for (int num : list) {
            if (num == pivot)
                equal.add(num);
            else if (num > pivot)
                bigger.add(num);
            else
                smaller.add(num);
        }
        // 如果bigger中的元素大于等于k个，说明答案在bigger中
        if (bigger.size() >= k)
            return findKth(bigger, k);
        // 否则，如果bigger+equal中的元素大于等于k个，说明答案在equal中
        else if (bigger.size() + equal.size() >= k)
            return pivot;
        // 否则，答案在smaller中
        else
            return findKth(smaller, k - bigger.size() - equal.size());
    }

    // 215. 数组中的第K个最大元素 另一种写法，更加高效
    public int findKthLargest_(int[] nums, int k) {
        return findKth_(nums, 0, nums.length - 1, k);
    }

    private int findKth_(int[] nums, int left, int right, int k) {
        // 随机选择pivot
        Random random = new Random();
        int pivot = nums[left + random.nextInt(right - left + 1)];
        // 三路快排的形式
        int lt = left, ge = right;
        int i = left;
        while (i <= ge) {
            if (nums[i] < pivot) {
                swap(nums, i, lt);
                lt++;
                i++;
            } else if (nums[i] > pivot) {
                swap(nums, i, ge);
                ge--;
            } else {
                i++;
            }
        }
        int smaller = lt - left, equal = ge - lt + 1, bigger = right - ge;
        if (bigger >= k)
            return findKth_(nums, ge + 1, right, k);
        else if (bigger + equal >= k)
            return pivot;
        else
            return findKth_(nums, left, lt - 1, k - bigger - equal);
    }

    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }

    // 295. 数据流的中位数
    class MedianFinder {

        PriorityQueue<Integer> large, small;

        public MedianFinder() {
            // 小顶堆
            large = new PriorityQueue<>((o1, o2) -> Integer.compare(o1, o2));
            // 大顶堆
            small = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1));
        }

        public void addNum(int num) {
            // 保证large的元素比small多一个，或者相同
            if (large.isEmpty()) {
                large.offer(num);
                return;
            }
            if (num > large.peek()) {
                large.offer(num);
                if (large.size() - small.size() >= 2)
                    small.offer(large.poll());
            } else {
                small.offer(num);
                if (small.size() > large.size())
                    large.offer(small.poll());
            }
        }

        public double findMedian() {
            if (large.size() == small.size())
                return (large.peek() + small.peek()) / 2.0;
            else
                return (double) large.peek();
        }
    }
}
